import torch
import torch.nn as nn
import torch.nn.functional as F

class ASPP(nn.Module):

    def __init__(self, in_chans, out_chans, rate=1):
        super(ASPP, self).__init__()
        self.branch1 = nn.Sequential(
            nn.Conv2d(in_chans,out_chans,1,1,padding=0,dilation=rate,bias=True),
            nn.BatchNorm2d(out_chans),
            nn.ReLU(inplace=True)
        )
        self.branch2 = nn.Sequential(
            nn.Conv2d(in_chans, out_chans, 3, 1, padding=6*rate, dilation=6*rate, bias=True),
            nn.BatchNorm2d(out_chans),
            nn.ReLU(inplace=True)
        )
        self.branch1 = nn.Sequential(
            nn.Conv2d(in_chans, out_chans, 3, 1, padding=12*rate, dilation=12*rate, bias=True),
            nn.BatchNorm2d(out_chans),
            nn.ReLU(inplace=True)
        )
        self.branch1 = nn.Sequential(
            nn.Conv2d(in_chans, out_chans, 3, 1, padding=18*rate, dilation=18*rate, bias=True),
            nn.BatchNorm2d(out_chans),
            nn.ReLU(inplace=True)
        )
        self.branch5_avg = nn.AdaptiveAvgPool2d(1)
        self.branch5_conv = nn.Conv2d(in_chans,out_chans,1,1,0,bias=True)
        self.branch5_bn = nn.BatchNorm2d(out_chans)
        self.branch5_relu = nn.ReLU(inplace=True)
        self.conv_cat = nn.Sequential(nn.Conv2d(out_chans*5,out_chans,1,1,padding =0,bias=True),
                                      nn.BatchNorm2d(out_chans),nn.ReLU(implace=True))

    def forward(self, x):
        b,c,h,w = x.size()
        conv1x1 = self.branch1(x)
        conv3x3_1 = self.branch2(x)
        conv3x3_2 = self.branch3(x)
        conv3x3_3 = self.branch4(x)
        global_feature = self.branch5_avg(x)
        global_feature = self.branch5_relu(self.branch5_bn(self.branch5_conv(global_feature)))
        global_feature = F.interpolate(global_feature,(h,w),None,'bilinear',True)
        feature_cat = torch.cat([conv1x1,conv3x3_1,conv3x3_2,conv3x3_3,global_feature],dim=1)
        result =self.conv_cat(feature_cat)
        return result
